// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/dwio/common/reader.h>

namespace kumo::pollux::dwio::common {
    using namespace pollux::common;

    VectorPtr RowReader::projectColumns(
        const VectorPtr &input,
        const ScanSpec &spec,
        const Mutation *mutation) {
        auto *inputRow = input->as<RowVector>();
        POLLUX_CHECK_NOT_NULL(inputRow);
        auto &inputRowType = input->type()->as_row();
        column_index_t numColumns = 0;
        for (auto &childSpec: spec.children()) {
            numColumns = std::max(numColumns, childSpec->channel() + 1);
        }
        std::vector<std::string> names(numColumns);
        std::vector<TypePtr> types(numColumns);
        std::vector<VectorPtr> children(numColumns);
        std::vector<uint64_t> passed(bits::nwords(input->size()), -1);
        if (mutation) {
            if (mutation->deletedRows) {
                bits::andWithNegatedBits(
                    passed.data(), mutation->deletedRows, 0, input->size());
            }
            if (mutation->randomSkip) {
                bits::forEachSetBit(passed.data(), 0, input->size(), [&](auto i) {
                    if (!mutation->randomSkip->testOne()) {
                        bits::clearBit(passed.data(), i);
                    }
                });
            }
        }
        for (auto &childSpec: spec.children()) {
            POLLUX_CHECK_NULL(childSpec->deltaUpdate());
            VectorPtr child;
            if (childSpec->isConstant()) {
                child = BaseVector::wrap_in_constant(
                    input->size(), 0, childSpec->constantValue());
            } else {
                child =
                        inputRow->childAt(inputRowType.getChildIdx(childSpec->fieldName()));
                childSpec->applyFilter(*child, passed.data());
            }
            if (!childSpec->projectOut()) {
                continue;
            }
            auto i = childSpec->channel();
            names[i] = childSpec->fieldName();
            types[i] = child->type();
            children[i] = std::move(child);
        }
        auto rowType = ROW(std::move(names), std::move(types));
        auto size = bits::countBits(passed.data(), 0, input->size());
        if (size == 0) {
            return RowVector::createEmpty(rowType, input->pool());
        }
        if (size < input->size()) {
            auto indices = allocate_indices(size, input->pool());
            auto *rawIndices = indices->asMutable<vector_size_t>();
            vector_size_t j = 0;
            bits::forEachSetBit(
                passed.data(), 0, input->size(), [&](auto i) { rawIndices[j++] = i; });
            for (auto &child: children) {
                child->disable_memo();
                child = BaseVector::wrap_in_dictionary(
                    nullptr, indices, size, std::move(child));
            }
        }
        return std::make_shared<RowVector>(
            input->pool(), rowType, nullptr, size, std::move(children));
    }

    namespace {
        void fillRowNumberVector(
            VectorPtr &rowNumVector,
            uint64_t previousRow,
            uint64_t rowsToRead,
            const dwio::common::SelectiveColumnReader *columnReader,
            VectorPtr &result) {
            FlatVector<int64_t> *flatRowNum{nullptr};
            if (rowNumVector && BaseVector::is_vector_writable(rowNumVector)) {
                flatRowNum = rowNumVector->as_flat_vector<int64_t>();
            }
            if (flatRowNum) {
                flatRowNum->clear_all_nulls();
                flatRowNum->resize(result->size());
            } else {
                rowNumVector = std::make_shared<FlatVector<int64_t> >(
                    result->pool(),
                    BIGINT(),
                    nullptr,
                    result->size(),
                    AlignedBuffer::allocate<int64_t>(result->size(), result->pool()),
                    std::vector<BufferPtr>());
                flatRowNum = rowNumVector->as_unchecked<FlatVector<int64_t> >();
            }
            auto *rawRowNum = flatRowNum->mutableRawValues();
            const auto rowOffsets = columnReader->outputRows();
            POLLUX_DCHECK_EQ(rowOffsets.size(), result->size());
            for (int i = 0; i < rowOffsets.size(); ++i) {
                rawRowNum[i] = previousRow + rowOffsets[i];
            }
        }
    } // namespace

    void RowReader::readWithRowNumber(
        std::unique_ptr<dwio::common::SelectiveColumnReader> &columnReader,
        const dwio::common::RowReaderOptions &options,
        uint64_t previousRow,
        uint64_t rowsToRead,
        const dwio::common::Mutation *mutation,
        VectorPtr &result) {
        const auto &rowNumberColumnInfo = options.rowNumberColumnInfo();
        POLLUX_CHECK(rowNumberColumnInfo.has_value());
        const auto rowNumberColumnIndex = rowNumberColumnInfo->insertPosition;
        const auto &rowNumberColumnName = rowNumberColumnInfo->name;
        column_index_t numChildren{0};
        for (auto &column: options.scanSpec()->children()) {
            if (column->keepValues()) {
                ++numChildren;
            }
        }
        POLLUX_CHECK_LE(rowNumberColumnIndex, numChildren);
        auto *row_vector = result->as_unchecked<RowVector>();
        VectorPtr rowNumVector;
        if (row_vector->childrenSize() != numChildren) {
            POLLUX_CHECK_EQ(row_vector->childrenSize(), numChildren + 1);
            rowNumVector = row_vector->childAt(rowNumberColumnIndex);
            const auto &rowType = row_vector->type()->as_row();
            auto names = rowType.names();
            auto types = rowType.children();
            auto children = row_vector->children();
            POLLUX_DCHECK(!names.empty() && !types.empty() && !children.empty());
            names.erase(names.begin() + rowNumberColumnIndex);
            types.erase(types.begin() + rowNumberColumnIndex);
            children.erase(children.begin() + rowNumberColumnIndex);
            result = std::make_shared<RowVector>(
                row_vector->pool(),
                ROW(std::move(names), std::move(types)),
                row_vector->nulls(),
                row_vector->size(),
                std::move(children));
        }
        columnReader->next(rowsToRead, result, mutation);
        fillRowNumberVector(
            rowNumVector, previousRow, rowsToRead, columnReader.get(), result);
        row_vector = result->as_unchecked<RowVector>();
        auto &rowType = row_vector->type()->as_row();
        auto names = rowType.names();
        auto types = rowType.children();
        auto children = row_vector->children();
        names.insert(names.begin() + rowNumberColumnIndex, rowNumberColumnName);
        types.insert(types.begin() + rowNumberColumnIndex, BIGINT());
        children.insert(children.begin() + rowNumberColumnIndex, rowNumVector);
        result = std::make_shared<RowVector>(
            row_vector->pool(),
            ROW(std::move(names), std::move(types)),
            row_vector->nulls(),
            row_vector->size(),
            std::move(children));
    }

    namespace {
        void logTypeInequality(
            const Type &fileType,
            const Type &tableType,
            const std::string &fileFieldName,
            const std::string &tableFieldName) {
            VKLOG(1) << "Type of the File field '" << fileFieldName
          << "' does not match the type of the Table field '" << tableFieldName
          << "': [" << fileType.toString() << "] vs [" << tableType.toString()
          << "]";
        }

        // Forward declaration for general type tree recursion function.
        TypePtr updateColumnNamesImpl(
            const TypePtr &fileType,
            const TypePtr &tableType,
            const std::string &fileFieldName,
            const std::string &tableFieldName);

        // Non-primitive type tree recursion function.
        template<typename T>
        TypePtr updateColumnNamesImpl(
            const TypePtr &fileType,
            const TypePtr &tableType) {
            const auto fileRowType = std::dynamic_pointer_cast<const T>(fileType);
            const auto tableRowType = std::dynamic_pointer_cast<const T>(tableType);

            std::vector<std::string> newFileFieldNames;
            newFileFieldNames.reserve(fileRowType->size());
            std::vector<TypePtr> newFileFieldTypes;
            newFileFieldTypes.reserve(fileRowType->size());

            for (auto childIdx = 0; childIdx < tableRowType->size(); ++childIdx) {
                if (childIdx >= fileRowType->size()) {
                    break;
                }

                newFileFieldTypes.push_back(updateColumnNamesImpl(
                    fileRowType->childAt(childIdx),
                    tableRowType->childAt(childIdx),
                    fileRowType->nameOf(childIdx),
                    tableRowType->nameOf(childIdx)));

                newFileFieldNames.push_back(tableRowType->nameOf(childIdx));
            }

            for (auto childIdx = tableRowType->size(); childIdx < fileRowType->size();
                 ++childIdx) {
                newFileFieldTypes.push_back(fileRowType->childAt(childIdx));
                newFileFieldNames.push_back(fileRowType->nameOf(childIdx));
            }

            return std::make_shared<const T>(
                std::move(newFileFieldNames), std::move(newFileFieldTypes));
        }

        // General type tree recursion function.
        TypePtr updateColumnNamesImpl(
            const TypePtr &fileType,
            const TypePtr &tableType,
            const std::string &fileFieldName,
            const std::string &tableFieldName) {
            // Check type kind equality. If not equal, no point to continue down the
            // tree.
            if (fileType->kind() != tableType->kind()) {
                logTypeInequality(*fileType, *tableType, fileFieldName, tableFieldName);
                return fileType;
            }

            // For leaf types we return type as is.
            if (fileType->isPrimitiveType()) {
                return fileType;
            }

            if (fileType->is_row()) {
                return updateColumnNamesImpl<RowType>(fileType, tableType);
            }

            if (fileType->is_map()) {
                return updateColumnNamesImpl<MapType>(fileType, tableType);
            }

            if (fileType->is_array()) {
                return updateColumnNamesImpl<ArrayType>(fileType, tableType);
            }

            // We should not be here.
            VKLOG(1) << "Unexpected table type during column names update for File field '"
          << fileFieldName << "': [" << fileType->toString() << "]";
            return fileType;
        }
    } // namespace

    TypePtr Reader::updateColumnNames(
        const TypePtr &fileType,
        const TypePtr &tableType) {
        return updateColumnNamesImpl(fileType, tableType, "", "");
    }
} // namespace kumo::pollux::dwio::common
